import matplotlib.pyplot as plt
import torch
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks
from data_convert import get_transform

# image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png")
image = read_image("/home/autox/zhangcong/zhangcong.png")
eval_transform = get_transform(train=False)

print(f"image.shape:{image.shape}")

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = torch.load("trainmodel.pth", weights_only=False)

model.eval()
with torch.no_grad():
    x = eval_transform(image)
    # convert RGBA -> RGB and move to device
    x = x[:3, ...].to(device)
    predictions = model([x, ])
    print(f"predictions:{predictions}")
    pred = predictions[0]

image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8)
image = image[:3, ...]
pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])]
pred_boxes = pred["boxes"].long()
print(pred_boxes.shape)
output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red")

print(pred["masks"].shape)
masks = (pred["masks"] > 0.7).squeeze(1)
print(masks.shape)
output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue")

plt.figure(figsize=(12, 12))
# plt.subplot(121)
plt.imshow(output_image.permute(1, 2, 0))
# plt.subplot(122)
# plt.imshow(masks.to("cpu").permute(1, 2, 0))
plt.show()